# -*- coding: utf-8 -*-
# created on 2016/5/10

from mathsolver.functions.base import *
from sympy import Abs, sqrt, Rational
from mathsolver.functions.sympy_utils import get_all_child, out_sqrt
from mathsolver.functions.fangcheng.budengshi import FangChenBuDengShi002, FangChenBuDengShi001


class HuaJian(BaseFunction):
    """
    HuaJian().solver(BasePoly("x**100+x**99+x**98+x**97+x**96"))
    HuaJian().solver(BasePoly("(a)**(3)/(a)*(1)/(a)"))
    HuaJian().solver(BasePoly("(y-x)**(2009)"))
    HuaJian().solver(BasePoly("5*(3*(a)**(2)*b-a*(b)**(2))-(a*(b)**(2)+3*(a)**(2)*b)+2*a*(b)**(2)"))
    HuaJian().solver(BasePoly("(a+3)**2-a*(4+a)"))
    HuaJian().solver(BasePoly("x**2 + 1/x**2"))
    HuaJian().solver(BasePoly("x"))
    """
    def solver(self, *args):
        expr = args[0].sympify()
        assert len(expr.free_symbols) > 0, "need at least one free symbol"

        # x^2007
        if get_all_child(expr, lambda arg: arg.is_Pow and (arg.args[1].free_symbols or arg.args[1] > 3)):
            self.output.append(BasePoly(expr))
            return self

        if not expr.is_Add and not expr.is_Mul:
            self.output.append(BasePoly(expr))
            return self

        self.steps.append(["", "%s" % new_latex(expr)])
        input_expr = expr

        simp = expr.expand().simplify()
        if simp != expr:
            self.steps.append(["", "= %s" % new_latex(simp)])
            expr = simp

        factor = expr.factor()
        if factor != expr:
            self.steps.append(["", "= %s" % new_latex(factor)])
            expr = factor

        cancel = expr.cancel()
        if str(cancel) != str(expr):
            self.steps.append(["", "= %s" % new_latex(cancel)])
            expr = cancel

        if len(self.steps) <= 1:
            self.steps = []
        elif new_latex(input_expr) == new_latex(expr):
            self.steps = []
        else:
            self.label.add("代数式化简")

        self.output.append(BasePoly(expr))
        return self


class HuaJianWithIneqs(BaseFunction):
    """
    HuaJianWithIneqs().solver(BaseIneqs([["1", '<', "x"], ["x", '<', "4"]]), BasePoly("sqrt((x-4)**2)+sqrt((x-1)**2)"))
    HuaJianWithIneqs().solver(BaseIneqs([["1", '<', "x"], ["x", '<', "4"]]), BasePoly("sqrt((x-4)**2*x**2)"))
    HuaJianWithIneqs().solver(BaseIneqs([["x", '<', "0"], ["y", '>', "0"]]), BasePoly("sqrt((x)**2*y)"))
    HuaJianWithIneqs().solver(BaseIneqs([["x", '<', "0"], ["y", '>', "0"]]), BasePoly("sqrt((x)**2)*sqrt(y*x**2)"))
    HuaJianWithIneqs().solver(BaseIneqs([["a", '<', "-b"], ["a/b", '>', "0"]]), BasePoly("Abs(a)-Abs(b)+Abs(a+b)+Abs(ab)"))
    """
    def solver(self, *args):
        ineqs = args[0]
        expr = args[1].sympify()
        expr_symbols = expr.atoms(Symbol)
        if len(expr_symbols) == 1:  # 表达式中只有一个变量的情形
            huajian_expr = self.solver_one_symbol(args[0], args[1])
        elif len(expr_symbols) >= 2 and (expr.is_Pow or expr.is_Mul):  # 表达式中有多个变量且单项式的情形
            if expr.is_Mul:
                huajian_expr = 1
                for i in range(len(expr.args)):
                    huajian_expr = huajian_expr * self.solver_symbols_one_item(args[0], BasePoly(expr.args[i]))
            elif expr.is_Pow:
                huajian_expr = self.solver_symbols_one_item(args[0], args[1])
        elif len(expr_symbols) >= 2 and expr.is_Add:  # 表达式中有多个变量且多项式情形
            huajian_expr = self.solver_symbols_items(args[0], args[1])
        else:
            raise ValueError("unexpected situation")
        self.steps.append(['因为%s，' % new_latex(ineqs.printing()), '故原式化简后的结果是：%s' % new_latex(huajian_expr)])
        self.output.append(BaseValue(huajian_expr))
        return self

    def huajian_abs(self, expr, symbol, symbol_interval):
        """
        化简表达式中指定的变量，将指定变量的绝对值去除,例：如果x<0,y<0，那么abs(x)*sqrt(y)->-x*sqrt(y)
        :param expr: 目标表达式
        :param symbol: 指定的变量名
        :param symbol_interval: 指定变量的取值范围
        :return: 化简后的表达式
        """
        huajian_expr = 1
        if expr.is_Mul:  # 当表达式是乘号连接的情况，例：abs(x)*sqrt(y)
            for mul_item in expr.args:
                if mul_item == -1:
                    huajian_expr = huajian_expr * mul_item
                    continue
                if isinstance(mul_item, Abs):
                    out_abs_expr = self.out_abs(mul_item, symbol_interval, symbol)
                    if mul_item != out_abs_expr:
                        huajian_expr = huajian_expr * self.out_abs(mul_item, symbol_interval, symbol)
                elif symbol == mul_item.free_symbols.pop():
                    huajian_expr = huajian_expr * mul_item
        elif isinstance(expr, Abs):  # 当表达式是abs的情况
            out_abs_expr = self.out_abs(expr, symbol_interval, symbol)
            if expr != out_abs_expr:
                huajian_expr = huajian_expr * self.out_abs(expr, symbol_interval, symbol)
        else:
            raise ValueError('unexpected situation')
        return huajian_expr

    def fenjie_sqrt(self, expr):
        """
        将根号内的表达式进行分解，例：sqrt(x**2*y**2)->sqrt(x**2)*sqrt(y**2)
        :param expr: 需要分解的表达式
        :return: 分解后的表达式
        """
        fenjie_expr = 1
        if expr.is_Pow and expr.args[0].is_Mul:
            for item in expr.args[0].args:
                fenjie_expr = fenjie_expr * sqrt(item)
            self.steps.append(['依题意，', '%s=%s' % (new_latex(expr), new_latex(fenjie_expr))])
        elif isinstance(expr.args[0], Symbol):
            fenjie_expr = expr
        elif expr.args[0].is_Pow:
            fenjie_expr = expr
        else:
            raise ValueError('unexpected situation')
        return fenjie_expr

    @staticmethod
    def out_sqrt(sqrt_expr):
        """
        去除表达式中的二次项根号，结果以绝对值形式给出，例：sqrt(x**2)->abs(x)
        :param sqrt_expr: 目标表达式
        :return: 移除二次根号后的表达式
        """
        if not sqrt_expr.is_Pow:  # 如果表达式不是Pow类型，直接返回原表达式
            return sqrt_expr
        elif sqrt_expr.args[1] == Rational(1, 2) and sqrt_expr.args[0].is_Pow and sqrt_expr.args[0].args[1] == 2:
            # 当根号内是平方表达式时
            return out_sqrt(sqrt_expr)
        elif isinstance(sqrt_expr.args[0], Symbol):  # 当根号内是一个字母时
            return sqrt_expr
        else:
            raise ValueError("unexpected situation!")

    @staticmethod
    def out_abs(expr, dingyiyu, symbol):
        """
        根据abs表达式中指定变量的取值范围，将指定变量去除绝对值，例：若x<0, abs(x)*sqrt(y)->-x*sqrt(y)
        :param expr: 表达式
        :param dingyiyu: 指定变量的取值范围
        :param symbol: 指定变量，需要移除绝对值的变量名
        :return: 被移除绝对值后的表达式
        """
        assert isinstance(expr, Abs)
        if symbol != expr.free_symbols.pop():
            return expr
        expr = expr.args[0]
        result_interval = FangChenBuDengShi001().solver(BaseIneq([expr, '>=', "0"]), BaseVariable(symbol)).output[
            0].interval
        if result_interval.intersect(dingyiyu) == dingyiyu:
            return expr
        else:
            result_interval = FangChenBuDengShi001().solver(BaseIneq([-expr, '>=', "0"]), BaseVariable(symbol)).output[0].interval
            if result_interval.intersect(dingyiyu) == dingyiyu:
                return -expr
            else:
                raise ValueError('This situation has not yet been considered!')

    def get_all_child_signs(self, expr, func):
        """
        提取表达式中的根号表达式以及对应的符号(正号和负号)
        :param expr: 目标表达式
        :param func: 判断表达式是二次根号的条件函数
        :return: 表达式与其相对应的符号字典的集合
        """
        childrens = []
        if func(expr):
            return [{expr: "+"}]
        if hasattr(expr, "args"):
            for arg in expr.args:
                if func(arg):
                    childrens.append({arg: "+"})
                else:
                    child = self.get_all_child_signs(arg, func)
                    if arg.is_Mul and arg.args[0].is_negative:  # 当表达式前面是负号，提取负号
                        child[0][arg.args[1]] = "-"
                    childrens.extend(child)
            return childrens
        else:
            return childrens

    @staticmethod
    def zhuanhua_expr(ineqs):
        """
        将不等式转化成统一形式，例如：a<-b -> a+b<0 或者 a/b>0 -> ab>0
        即将不等式转化成：不等号左边是表达式，右边是0的情况，且将除号转化成乘号
        :param ineqs: 不等式
        :return: 统一形式后的不等号
        """
        value = ineqs
        if value[0] != 0 and value[2] != 0:
            value0 = value[0] - value[2]
            value1 = value[1]
            value2 = S.Zero
            return [value0, value1, value2]
        elif value[0].is_Mul and value[2] == 0:
            if value[0].args[1].is_Pow and value[0].args[1].args[1] == -1:
                value0 = value[0] * value[0].args[1].args[0] * value[0].args[1].args[0]
                value1 = value[1]
                value2 = value[2]
            return [value0, value1, value2]
        return [value[0], value[1], value[2]]

    def get_dingyiyu(self, ineqs, expr_symbols):
        """
        获取特殊情形的定义域，例如：a+b<0,ab>0 -> a<0,b<0
        :param ineqs: 不等式组
        :param expr_symbols: 不等式组中的原子变量
        :return: 返回结果
        """
        assert len(ineqs.ineqs) == 2 and len(expr_symbols) == 2
        symbol0 = expr_symbols.pop()
        symbol1 = expr_symbols.pop()
        expr_symbols.add(symbol0)
        expr_symbols.add(symbol1)
        ineq_0 = self.zhuanhua_expr(ineqs.ineqs[0].sympify())
        ineq_1 = self.zhuanhua_expr(ineqs.ineqs[1].sympify())
        if (ineq_0 == [symbol0 + symbol1, '<', 0] or ineq_1 == [symbol0 + symbol1, '<', 0]) and \
                (ineq_0 == [symbol0 * symbol1, '>', 0] or ineq_1 == [symbol0 * symbol1, '>', 0]):
            return BaseIneqs([[symbol0, '<', "0"], [symbol1, '<', "0"]])

    def solver_one_symbol(self, *args):
        """
        表达式中只有一个变量的情形
        :param args:
        :return:
        """
        ineqs = args[0]
        expr = args[1].sympify()
        symbol = expr.free_symbols.pop()
        dingyiyu = FangChenBuDengShi002().solver(ineqs, BaseVariable(symbol)).output[0].value[symbol]
        self.steps.append(['因为%s，' % ineqs.printing(), '所以%s的范围在%s范围内' % (new_latex(symbol), new_latex(dingyiyu))])
        children = self.get_all_child_signs(expr, lambda arg: arg.is_Pow and arg.args[1].is_Number)
        huajian_expr = sympify(0)
        for sqrt_item in children:
            sqrt_item_key = list(sqrt_item.keys())[0]
            fenjie_expr = self.fenjie_sqrt(sqrt_item_key)  # 分解sqrt，例：sqrt(x**2*y**2)->sqrt(x**2)*sqrt(y**2)
            if fenjie_expr.is_Mul:  # 当表达式是由乘号相连的时候
                out_sqrt_expr = 1
                for mul_item in fenjie_expr.args:
                    out_sqrt_expr = out_sqrt_expr * self.out_sqrt(mul_item)
            elif fenjie_expr.is_Pow:
                out_sqrt_expr = self.out_sqrt(fenjie_expr)  # 开根号，结果以绝对值形式给出，例：sqrt(x**2*y)->abs(x)*sqrt(y)
            # 根据定义域化简绝对值，例：若x<0, abs(x)*sqrt(y)->-x*sqrt(y)
            out_abs_expr = self.huajian_abs(out_sqrt_expr, symbol, dingyiyu)
            sign = sqrt_item[sqrt_item_key]
            self.steps.append(['表达式%s' % new_latex(sqrt_item_key), '化简后的结果为：%s' % new_latex(out_abs_expr)])
            if sign == "+":
                huajian_expr = huajian_expr + out_abs_expr
            else:
                huajian_expr = huajian_expr - out_abs_expr
        return huajian_expr

    def solver_symbols_one_item(self, *args):
        """
        表达式中有两个或者多个变量，且表达式是单项式
        :param args:
        :return:
        """
        ineqs = args[0]
        expr = args[1].sympify()
        expr_symbols = expr.atoms(Symbol)  # 获取表达式中的原子变量，如：a,b
        if expr.is_Pow and expr.args[1] == Rational(1, 2):  # 处理表达式是根号式的情况
            huajian_expr = 1
            fenjie_expr = self.fenjie_sqrt(expr)  # 将根号式分解，如：sqrt(a**2*b) -> sqrt(a**2)*sqrt(b)
            if fenjie_expr.is_Mul:  # 当表达式是由乘号相连的时候
                out_sqrt_expr = 1
                for mul_item in fenjie_expr.args:
                    out_sqrt_expr = out_sqrt_expr * self.out_sqrt(mul_item)
            elif fenjie_expr.is_Pow:
                out_sqrt_expr = self.out_sqrt(fenjie_expr)  # 开根号，结果以绝对值形式给出，例：sqrt(x**2*y)->abs(x)*sqrt(y)
            for i in range(len(expr_symbols)):  # 遍历表达式中的每一个原子式，并将对应的根号式化简
                assert len(expr_symbols) > 0
                symbol = expr_symbols.pop()
                symbol_interval = FangChenBuDengShi001().solver(ineqs.ineqs[i], BaseVariable(symbol)).output[
                    0].interval
                huajian_expr = huajian_expr * self.huajian_abs(out_sqrt_expr, symbol, symbol_interval)
        elif isinstance(expr, Symbol):
            huajian_expr = expr
        else:
            raise ValueError("unexpected situation")
        return huajian_expr

    def solver_symbols_items(self, *args):
        """
        表达式中有多个变量的情形，且表达式是多项式时
        :param args:
        :return:
        """
        ineqs = args[0]
        expr = args[1].sympify()
        expr_symbols = expr.atoms(Symbol)
        symbol_set = set()
        for item in expr_symbols:
            if len(item.name) == 1:
                symbol_set.add(item)
        expr_symbols = symbol_set
        huajian_expr = 0
        if expr.is_Add and len(expr_symbols) == 2:
            ineqs = self.get_dingyiyu(ineqs, expr_symbols)
            symbol_interval = FangChenBuDengShi001().solver(ineqs.ineqs[0], BaseVariable(expr_symbols.pop())).output[
                0].interval
            for abs_item in expr.args:
                abs_symbol = abs_item.atoms(Symbol).pop()
                if len(abs_item.atoms(Symbol)) == 1 and len(abs_symbol.name) == 1:
                    result_expr = self.huajian_abs(abs_item, abs_item.free_symbols.pop(), symbol_interval)
                elif len(abs_item.atoms(Symbol)) == 1 and len(abs_symbol.name) == 2:
                    result_expr = abs_item.args[0]
                elif len(abs_item.atoms(Symbol)) == 2:
                    if symbol_interval.intersect(Interval(-S.Infinity, 0, True, False)) == symbol_interval:
                        result_expr = -abs_item.args[0]
                    elif symbol_interval.intersect(Interval(0, S.Infinity, False, True)) == symbol_interval:
                        result_expr = abs_item.args[0]
                if result_expr != 1:
                    huajian_expr = huajian_expr + result_expr
        return huajian_expr


if __name__ == '__main__':
    pass
